#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 2018/1/24 15:03
@desc: 
"""
import os

import cv2
import numpy as np
from skimage import transform

from utils.basic import get_file_name
from utils.list import create_file_list


def get_image_list(root, min_number=0, max_number=100000):
    """
    从已经分割好的图像的文件目录获取图像列表，只用于读取小图

    Parameters
    ----------
    root
    min_number
    max_number

    Returns
    -------

    """
    def _filter_number_of_person(file_path, lists, params):
        if not file_path.endswith('.jpg'):
            return lists
        _root = os.path.dirname(file_path)
        name = get_file_name(file_path)
        if os.path.exists(os.path.join(_root, name + '.txt')):
            with open(os.path.join(_root, name + '.txt'), 'r') as file:
                if params['min'] < len(file.readlines()) < params['max']:
                    lists[0].append(os.path.join(_root, name + '.txt') + params['split'] + file_path)
        return lists

    file_list, _ = create_file_list(root, filtering=_filter_number_of_person,
                                    params={'min': min_number, 'max': max_number,
                                            'split': '&!&'})
    return file_list


def convert_density_map(origin_map, shape=(320, 320), resize_shape=None):
    """
    将density map转化为其他shape

    Parameters
    ----------
    origin_map: 原始density map
    shape: 输出的density map尺寸
    resize_shape: 图像实际被resize后的大小

    Returns
    -------
    目标尺寸的density map
    """
    if resize_shape is None:
        resize_shape = shape
    else:
        assert resize_shape[0] <= shape[0] and resize_shape[1] <= shape[1], \
            'shape is smaller than resize shape !'
    pad_h, pad_w = (shape[0] - resize_shape[0]) // 2, (shape[1] - resize_shape[1]) // 2
    buffer_map = origin_map[pad_h:(pad_h + resize_shape[0]), pad_w:(pad_w + resize_shape[1])]
    summary = np.sum(buffer_map)
    if resize_shape[0] == shape[0] and resize_shape[1] == shape[1]:
        density_map = buffer_map
    else:
        density_map = transform.resize(buffer_map, shape, mode='symmetric', preserve_range=True)
    density_map = summary * density_map / np.sum(density_map)
    return density_map


def preprocess_image(image, mean=(80.197326, 75.26667, 72.822747),
                     shape=(256, 256), resize_shape=None):
    """
    对要检测的图像进行预处理， 包括resize、substract mean

    Parameters
    ----------
    image: 输入图像
    shape: 输出图像的shape
    mean: 各个channel的均值，依次为R、G、B
    resize_shape: 图像实际被resize后的大小

    Returns
    -------
    处理后的图像

    """
    if resize_shape is None:
        resize_shape = shape
    else:
        assert resize_shape[0] <= shape[0] and resize_shape[1] <= shape[1], \
            'shape is smaller than resize shape !'
    img = image.astype(np.float32)
    img = cv2.resize(img, resize_shape)
    img[:, :, 0] -= mean[0]
    img[:, :, 1] -= mean[1]
    img[:, :, 2] -= mean[2]
    pad_h, pad_w = (shape[0] - resize_shape[0]) // 2, (shape[1] - resize_shape[1]) // 2
    bias_h, bias_w = (shape[0] - resize_shape[0]) % 2, (shape[1] - resize_shape[1]) % 2
    img = np.pad(img, ((pad_h, pad_h + bias_h), (pad_w, pad_w + bias_w), (0, 0)),
                 'constant', constant_values=0)
    return img


def rotate_image(angle, image):
    """
    逆时针旋转一个图像 angle 度

    Parameters
    ----------
    angle: 要旋转的角度
    image: 原始图像

    Returns
    -------
    img_buf: 从旋转后的图像的中心切割出来与原始图像相同shape的图像
    offset: img_buf 相对于旋转后的图像的顶点（左上角）的起始坐标

    """
    shape = image.shape
    arc = angle * np.pi / 180
    post_size = (np.ceil(abs(shape[0] * np.cos(arc)) + abs(shape[1] * np.sin(arc))).astype(np.int32),
                 np.ceil(abs(shape[0] * np.sin(arc)) + abs(shape[1] * np.cos(arc))).astype(np.int32))
    offset = list(map(lambda x, y: (x - y) // 2, post_size, shape))
    if 0 == angle:
        img_buf = image.copy()
        img_buf = img_buf.astype(np.float32)
    else:
        img_buf = transform.rotate(image, angle, resize=True, preserve_range=True)
        img_buf = img_buf[offset[0]:(offset[0] + shape[0]), offset[1]:(offset[1] + shape[1]), :]
    return offset, img_buf
